
import unittest
import diabind
import testdata
import os

class SessionTest( unittest.TestCase ):

    def __get_func(self):
        if not hasattr(self, "__func"):
            glob = testdata.loaded_session.get_globalScope()
            enum = glob.findChildren( diabind.SymTagFunction,
                                      "DllMain",
                                      diabind.nsfUndecoratedName )
            self.assertEqual(len(enum), 1)
            self.__func = enum[0]
        return self.__func

    def test_loadAddress(self):
        loadAddr = testdata.loaded_session.get_loadAddress() + 0x10000
        testdata.loaded_session.put_loadAddress(loadAddr)
        self.assertEqual(loadAddr, testdata.loaded_session.get_loadAddress())

    def test_findSymbolByAddr(self):
        func = testdata.loaded_session.findSymbolByAddr( self.__get_func().get_addressSection(),
                                                         self.__get_func().get_addressOffset(),
                                                         diabind.SymTagFunction )
        self.assertEqual(self.__get_func().get_symIndexId(), func.get_symIndexId())

    def __test_findSymbolEx(self, method, addr):
        find_res = method(addr, diabind.SymTagFunction)
        self.assertEqual(self.__get_func().get_symIndexId(), find_res[0].get_symIndexId())
        self.assertEqual(0, find_res[1])

        find_res = method(addr + 2, diabind.SymTagFunction)
        self.assertEqual(self.__get_func().get_symIndexId(), find_res[0].get_symIndexId())
        self.assertEqual(2, find_res[1])

    def test_findSymbolByRVAEx(self):
        self.__test_findSymbolEx( testdata.loaded_session.findSymbolByRVAEx,
                                  self.__get_func().get_relativeVirtualAddress() )

    def test_findSymbolByVAEx(self):
        self.__test_findSymbolEx( testdata.loaded_session.findSymbolByVAEx,
                                  self.__get_func().get_virtualAddress() )

    def test_symbolById(self):
        func = testdata.loaded_session.symbolById(self.__get_func().get_symIndexId())
        self.assertEqual(self.__get_func().get_symIndexId(), func.get_symIndexId())

    def __test_findLines(self, method, args):
        enum_lines = method( *(args + [1,]) )
        self.assertEqual(len(enum_lines), 1)

        enum_lines = method( *(args + [0x100,]) )
        self.assertTrue(len(enum_lines) > 1)

    def test_findLinesByAddr(self):
        self.__test_findLines( testdata.loaded_session.findLinesByAddr,
                               [ self.__get_func().get_addressSection(), 
                                 self.__get_func().get_addressOffset() ] )

    def test_findLinesByRVA(self):
        self.__test_findLines( testdata.loaded_session.findLinesByRVA,
                               [ self.__get_func().get_relativeVirtualAddress(), ] )

    def test_findLinesByVA(self):
        self.__test_findLines( testdata.loaded_session.findLinesByVA,
                               [ self.__get_func().get_virtualAddress(), ] )
